from torch.utils.data import Dataset
import torchvision.transforms as transforms
import numpy as np
from PIL import Image
import torch as t
import os
import numpy as numpy


class LabelUtil:
    # Introduction
    # 这个类的功能:
    #    Convert a tag image with 3 channels to a tag image with 1 channel.
    #    Each pixel value of the 1-channel image is an integer, each integer representing the corresponding RGB value in the original 3-channel tag, as well as a category
    #    将一张 彩色 Label 转换成 一个单通道的  图片数组，
    #    每个颜色用一个 整数代表，一个整数代表一个类别
    # Example:
    # ---------
    #  (0,0,0) ==> Background->(0), (0,255,127)==>Vegetation->(1), (220,20,60)==>Roads->(2),(255,255,0)==>Building->(3), (0,191,255)==> River->(4)
    #  color = [(0,0,0),(0,255,127),(220,20,60),(255,255,0),(0,191,255)]
    def __init__(self, color):
        # 创建一个Hash数组, 每一种颜色，给定一个下标值，这个下标所在的位置的值 是一个整数id，一个id，代表一个分类
        cm = np.zeros(256 ** 3)
        for cid, item in enumerate(color):
            cm[(item[0] * 256 + item[1]) * 256 + item[2]] = cid
        self.colormap = cm

    def __call__(self, pic):
        # three channels to one channels
        # RGB ---> I
        # 其中：
        #    RGB  ---> K --> I,  K = ( R * 256 + G) * 256 + B ----> 每个颜色有个 id， 参见 代码第18行。
        label = np.array(pic, dtype='uint32')
        cid = (label[:, :, 0] * 256 + label[:, :, 1]) * 256 + label[:, :, 2]
        #         Example:
        #         -------
        #           cid = ⇩                                                   return
        #             [[       0        0        0        0        0]         [[0 0 0 0 0]
        #              [       0    65407        0        0        0]          [0 1 0 0 0]
        #              [       0        0 14423100        0        0]   ===>   [0 0 2 0 0]
        #              [       0        0        0 16776960        0]          [0 0 0 3 0]
        #              [       0        0        0        0    49151]]         [0 0 0 0 4]]
        return np.array(self.colormap[cid], dtype='uint8')


class UserDataset(Dataset):
    def __init__(self, img_path: str, label_path: str, labelUtils: LabelUtil):
        super(UserDataset, self).__init__()
        # 初始化图片和标签全路径，数组
        self.imgslist = [os.path.join(img_path, item) for item in os.listdir(img_path)]
        self.labelslist = [os.path.join(label_path, item) for item in os.listdir(label_path)]
        self.transforms = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        self.labelUtils = labelUtils
        # 路径按名字，排个序
        self.imgslist.sort()
        self.labelslist.sort()

    def __getitem__(self, index):
        # open image with PIL
        img, label = Image.open(self.imgslist[index]), Image.open(self.labelslist[index])
        # 注意：
        #      label 和 image 如果尺寸相同，可往下走。若不同需要另外书写代码，用中心裁剪，变成尺寸相同的。
        #      如果label是单通道，可替换下面的RGB_Label_To_SingleChannel_To_Tensor 变为 t.from_numpy(label)
        #      下面处理的label 是3通道的彩色图
        labelx = Image.new('RGB', label.size)
        labelx.paste(label, (0, 0))

        img, label = self.transforms(img), t.from_numpy(self.labelUtils(label))

        sample = {'img': img, 'label': label}

        return sample

    def __len__(self):
        return len(self.imgslist)


# 创建数据集
#  1.初始化路径，以及标签
train_root_path = r'data/train'
label_root_path = r'data/train_labels'
color = [[0, 0, 0], [0, 255, 127], [255, 255, 0], [0, 191, 255], [220, 20, 60]]
# 2.
labelUtils = LabelUtil(color)
userDataset = UserDataset(train_root_path, label_root_path, labelUtils)
# print(userDataset[0]['img'].shape)
